#include "types.h"
#include "param.h"
#include "memlayout.h"
#include "mmu.h"
#include "proc.h"
#include "x86.h"
#include "spinlock.h"
#include "defs.h"
#include "shmflg.h"

#define framePerSegment 	500    //TODO: maybe you the params of number of pages
#define MaxSegments	500  //TODO: maybe you the params of number of pages
struct shm_segment
{
  int key;
  int shm_id;
  uint framesAmount;
  char* frames[framePerSegment];
  uint process_count;  
};

struct {
  struct spinlock lock;
  struct shm_segment table[MaxSegments];
} segments_table;

//help functions
int validatePages(struct shm_segment* seg,const void* va)
{
  //check the physical addresses at the page 
  //table correspond to the physical addresses at the segments table.
  char* pAddress;
  uint i;  
  for (i=0; i<seg->framesAmount; ++i)
  {
    pAddress= uva2ka(proc->pgdir,va);
    if (pAddress != seg->frames[i])
      return -1;
    va+=PGSIZE;
  }
  //every thing ok
  return 0;
  
}

struct shm_segment* findSegmentByAdress(char* pAddres)
{
  struct shm_segment* it;
  for (it=segments_table.table; it<&segments_table.table[MaxSegments]; ++it)
  {
    if (it->shm_id == -1)
      continue;
    
    if (it->frames[0] == pAddres)
      return it;
  }
  //not found
  return 0;  
}

void freePages(uint startAddres,int count)
{
    uint va  = startAddres;
    uint j;
    pte_t *pte;
    for (j=0; j<count; ++j)
    {	
      pte = walkpgdir(proc->pgdir, (void*)va, 0);
      if (pte == 0)
	panic("pte not supposed to be 0\n");
      *pte = 0;
      va+=PGSIZE;
    }
}

//return the start of virtual memory
void* attachMemoryToProc(struct shm_segment* seg,int perm)
{
  uint i;
  uint va = proc->sz;
  if(va+PGSIZE*seg->framesAmount >= KERNBASE)
  {
    cprintf("shmat: not enough memory to process %d to be attached\n", proc->pid);
    return (void*)-1;
  }
  for (i=0; i<seg->framesAmount; ++i)
  {    
    char* mem = seg->frames[i];    
    if (mappages(proc->pgdir, (char*)va, PGSIZE, v2p(mem), perm) <0)
    {
      cprintf("shmat: mappages failed \n");
      //free the pages we already attached 
      freePages(proc->sz,i);
      return (void*)-1;
    }
    va+=PGSIZE;
  }
  uint oldSize = proc->sz;
  proc->sz = va;
  return (void*)oldSize;
}
void freeMemory(struct shm_segment* seg)
{
  uint j;
  for (j=0; j<seg->framesAmount; ++j)
      kfree(seg->frames[j]);
  seg->framesAmount = 0;
}

struct shm_segment* findSegmentBySHMID(int shmID)
{
  if (shmID <0 || shmID >= MaxSegments)
    return 0;
  struct shm_segment* seg = &segments_table.table[shmID];
  if (seg->shm_id == -1)
    return 0;
  if (seg->shm_id != shmID)
  {    
    cprintf("shmID of segment at index %d not supposed to be %d\n",seg->shm_id,shmID);
    panic("shared memory");
  }
  return seg;
}

struct shm_segment* findSegmentByKey(int key)
{
  struct shm_segment* it;
  for (it=segments_table.table; it<&segments_table.table[MaxSegments]; ++it)
  {
    if (it->shm_id == -1)
      continue;
    
    if (it->key == key)
	return it;
  }
  //not found
  return 0;
}

struct shm_segment* allocateSegment(key)
{
  struct shm_segment* seg;
  uint i;
  for (i=0; i<MaxSegments; ++i)
  {
    seg = &segments_table.table[i];
    if (seg->shm_id == -1) //found unused segment
    {
      seg->shm_id = i;
      seg->key = key;
      seg->process_count =0;
      return seg;
    }
  }
  //the segment table is full
  return 0;
}

int allocateMemory(struct shm_segment* seg,uint size)
{
  size = PGROUNDUP(size);
  seg->framesAmount = size / PGSIZE;
  if (seg->framesAmount > framePerSegment) //not enough frame per segment
  {
    cprintf("shmget: process %d want to allocate % bytes, only %d pages allowed\n",proc->pid,size,framePerSegment);
    return -1;
  }
  uint i;
  for (i=0; i<seg->framesAmount; ++i)
  {    
    char* mem = kalloc();
    if(mem == 0){
      cprintf("shmat: shared memory out of memory\n");      
      uint j;
      //free what we already allocated
      for (j=0; j<i; ++j)
	kfree(seg->frames[j]);
      return -1;
    }
    memset(mem, 0, PGSIZE);
    seg->frames[i] = mem;
  }
  //everything ok
  return 0;
}

//system calls
void shm_init()
{
  initlock(&segments_table.lock,"shared memory");
  struct shm_segment* it;
  for (it=segments_table.table; it<&segments_table.table[MaxSegments]; ++it)
  {
    it->shm_id = -1;
  }
}
int shmget(int key, uint size, int shmflg)
{   
  struct shm_segment* seg;
  if (shmflg == GET)
  {
    acquire(&segments_table.lock);
    seg = findSegmentByKey(key);
    if (seg == 0)
    {
      release(&segments_table.lock);
      cprintf("shmget: segment with key %d not found \n",key);
      return -1;
    }
    release(&segments_table.lock);
    return seg->shm_id;
  }
  
  if (shmflg == CREAT)
  {
    acquire(&segments_table.lock);
    //search if segment for this key already exist
    seg = findSegmentByKey(key);
    if (seg!=0)
    {
      release(&segments_table.lock);
      cprintf("shmget: segement with key %d already exists\n",key);
      return -1;
    }
    
    //create segment
     seg = allocateSegment(key);
     if (seg == 0)
     {
       release(&segments_table.lock);
	cprintf("shmget: segment table is full \n");
	return -1; 
     }
     int res = allocateMemory(seg,size);
     if (res <0)
     {
       //make the segment unused
       seg->shm_id = -1;
       release(&segments_table.lock);
       return -1;
     }
     int shmid = seg->shm_id; 
     release(&segments_table.lock);
     return shmid;
  }
  
  //invalid flags
  cprintf("shmget got invalid flags\n");
  return -1;
}

int shmdel(int shmid)
{
  acquire(&segments_table.lock);
  struct shm_segment* seg =  findSegmentBySHMID(shmid);
  if (seg == 0)
  {
    release(&segments_table.lock);
    cprintf("shmdel: segment %d is not exist\n",shmid);
    return -1;
  }
  if (seg->process_count == 0)
  {
    uint frames = seg->framesAmount;
    freeMemory(seg);
    seg->shm_id = -1;
    release(&segments_table.lock);
    return frames;
  }
  else
  {
    release(&segments_table.lock);
    cprintf("ERROR: can't delete shared memory %d  bacause some process use it. \n",shmid);
    return -1;
  }
}

void* shmat(int shmid, int shmflg)
{
  if ((shmflg != SHM_RDONLY) && (shmflg != SHM_RDWR))
  {
    cprintf("ERROR: shmat - wrong flags\n");
    return (void*)-1;
  }
  acquire(&segments_table.lock);
  struct shm_segment* seg = findSegmentBySHMID(shmid);
  if (seg == 0) 
  {
    release(&segments_table.lock);
    cprintf("shmat: segment %d is not exist\n",shmid);
    return (void*)-1;
  }
  void* vadress;
  if (shmflg == SHM_RDONLY)
  {
    vadress = attachMemoryToProc(seg,PTE_U|PTE_P);
  }
  if (shmflg == SHM_RDWR)
  {
    vadress = attachMemoryToProc(seg,PTE_W|PTE_U|PTE_P);
  }
  seg->process_count++;
  release(&segments_table.lock);
  return vadress; 
}

int shmdt(const void *shmaddr)
{
  char* pAddress = uva2ka(proc->pgdir,shmaddr);
  if (pAddress == 0)
  {
    cprintf("shmdt: can't find page of virtual address %d\n",shmaddr);
    return -1;
  }
  acquire(&segments_table.lock);
 struct shm_segment* seg = findSegmentByAdress(pAddress); 
 if (seg == 0)
 {
   release(&segments_table.lock);
   cprintf("shmdt: no segment corresponding to virtual address %d\n",shmaddr);
   return -1;
 }
 int res = validatePages(seg,shmaddr);
 if (res<-1)
 {
   release(&segments_table.lock);
   cprintf("shmdt: segment frames not corresponding to process pages\n");
   return -1;
 }
 freePages((uint)shmaddr, seg->framesAmount); 
 seg->process_count--; 
 //sucsess 
 release(&segments_table.lock);
 return 0;
}